﻿#include <windows.h>
#include <iostream>
#include <Shlwapi.h>
#include <string>
#include "Encoding.h"

using std::wstring;

#define outline(value) std::cout << value << std::endl;
bool TryCreateChildProcess(PROCESS_INFORMATION& pi, const char* moduleFile, const char* arg);
void WaitForMutex(const char* mutexName);

#pragma comment(lib, "Shlwapi.lib")

/// <summary>
/// 主函数。
/// </summary>
/// <param name="argc">参数个数，必须固定4个。</param>
/// <param name="argv">第二个参数为互斥量名称，第二个参数为模块文件，第三个参数为模块执行命令参数。</param>
int main(int argc, char* argv[])
{
	if (argc != 5)
	{
		OutputDebugStringA("命令行参数错误！");
		return -1;
	}

	char* mutexName = argv[1];
	char* moduleFile = argv[2];
	char* args = argv[3];
	char* clientFile = argv[4];

	if (PathFileExistsA(moduleFile) == false)
	{
		OutputDebugStringA("模块文件不存在！");
		return -1;
	}

	PROCESS_INFORMATION serverInfo;

	if (TryCreateChildProcess(serverInfo, moduleFile, args) == false)
	{
		OutputDebugStringA("创建服务进程失败！");
		return -1;
	}

	WaitForMutex(mutexName);

	PROCESS_INFORMATION clientInfo;
	BOOL closeResult;

	// 尝试关闭 redis-server.
	outline("安全关闭 redis server ...");
	if (TryCreateChildProcess(clientInfo, clientFile, "shutdown") == false)
	{
		// 如果安全关闭失败，直接强制关闭。
		outline("强制关闭 redis 服务...");
		closeResult = TerminateProcess(serverInfo.hProcess, 0);

		if (closeResult == false)
		{
			outline("强制关闭 redis 服务失败。");
		}
	}
	else 
	{
		if (WaitForSingleObject(clientInfo.hProcess, 100) == WAIT_FAILED)
		{
			outline("强制关闭 redis 客户端。");
			closeResult = TerminateProcess(clientInfo.hProcess, 0);

			if (closeResult == false)
			{
				outline("强制关闭 redis 客户端失败。");
			}
		}

		CloseHandle(clientInfo.hProcess);
		CloseHandle(clientInfo.hThread);
	}

	CloseHandle(serverInfo.hProcess);
	CloseHandle(serverInfo.hThread);
	outline("redis shell exit.");
}

void WaitForMutex(const char* mutexName)
{
	outline("Open Mutex...");

	auto hMutex = OpenMutexA(MUTEX_ALL_ACCESS, false, mutexName);

	if (hMutex == NULL)
	{
		OutputDebugStringA("打开互斥锁失败！");
		return;
	}

	outline("Wait to get mutex.");
	auto waitStatus = WaitForSingleObject(hMutex, INFINITE);
	outline("I got the mutex..");

	if (waitStatus == WAIT_FAILED)
	{
		OutputDebugStringA("等待互斥锁失败！");
	}

	ReleaseMutex(hMutex);
	outline("Release the mutex..");
}

bool TryCreateChildProcess(PROCESS_INFORMATION& pi, const char* moduleFile, const char* args)
{
	STARTUPINFO si;

	ZeroMemory(&si, sizeof(si));
	si.cb = sizeof(si);
	ZeroMemory(&pi, sizeof(pi));

	auto wModuleFile = Encoding::ToWChar(moduleFile);
	auto wArg = Encoding::ToWChar(args);
	auto cmd = wModuleFile + L" " + wArg;

	// Start the child process. 
	if (!CreateProcess(NULL,   // No module name (use command line)
		(wchar_t*)cmd.c_str(),         // Command line
		NULL,           // Process handle not inheritable
		NULL,           // Thread handle not inheritable
		FALSE,          // Set handle inheritance to FALSE
		0,              // No creation flags
		NULL,           // Use parent's environment block
		NULL,           // Use parent's starting directory 
		&si,            // Pointer to STARTUPINFO structure
		&pi)           // Pointer to PROCESS_INFORMATION structure
		)
	{
		return false;
	}

	return true;
}